#include <buola/mat.h>
#include <buola/mat/prettyprint.h>
#include <buola/mat/markov.h>
#include <buola/mat/matrixops.h>
#include <buola/chart/cxychart.h>
#include <buola/chart/plots/clineplot.h>
#include <buola/chart/plots/cbarplot.h>

using namespace buola;

mat::CMat_d filter_hmm(const mat::CMat_d &pA,const mat::CMat_d &pB,const mat::CRow_d &pPi,const mat::CVec_i &pY)
{
    mat::CMat_d lPX=mat::zeros(pY.size(),pPi.size());
    
    lPX(0,nAll)=pB(nAll,pY(0)).T()**pPi;
    lPX(0,nAll)=lPX(0,nAll)/sum(lPX(0,nAll));
    
    for(int i=1;i<pY.size();i++)
    {
        lPX(i,nAll)=pB(nAll,pY(i)).T()**(lPX(i-1,nAll)*pA);
        lPX(i,nAll)=lPX(i,nAll)/sum(lPX(i,nAll));
    }
    
    return std::move(lPX);
}

int main(int pArgC,char** pArgs)
{
    buola_init(pArgC,pArgs);
    
    static const size_t lLen=1000000;

    mat::CMat_d lA{{0.9,0.1},{0.3,0.7}};
    mat::CRow_d lPi{0.5,0.5};

    mat::CVec_d lP=mat::linspace(50,0.0,1.0);
    mat::CVec_d lRes=mat::zeros(50);
    
    for(int i=0;i<lP.size();i++)
    {
        mat::CMat_d lB{{lP(i),1-lP(i)},{1-lP(i),lP(i)}};
    
        mat::CVec_i lX,lY;
    
        std::tie(lX,lY)=mat::generate_hmm(lLen,lA,lB,lPi);
        mat::CMat_d lPX=filter_hmm(lA,lB,lPi,lY);
    
        lRes(i)=count(indmax(rows(lPX))==lX)/double(lLen);
    }
    
    msg_info() << lRes << "\n";
    
    chart::PXYChart lChart=new chart::CXYChart;
    chart::PLinePlot lLine=new chart::CLinePlot;
    
    lLine->SetSeries(new chart::CSeries(lRes));
    
    lChart->AddPlot(lLine);
    lChart->Map(CSize_d(800,600));

    buola_main_loop();
    
    return buola_finish();
}
